{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_community.chat_models import ChatTongyi\n",
    "\n",
    "chat = ChatTongyi()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 阻塞模式"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content='勾股定理是古希腊数学家毕达哥拉斯发现的一个几何学基本定理，也被称为毕达哥拉斯定理。它描述了直角三角形中各边之间的关系：在一个直角三角形中，直角两边的平方和等于斜边的平方。用数学公式表示就是：\\n\\n如果直角三角形的两条直角边分别是a和b，斜边是c，那么有：\\na² + b² = c²\\n\\n这个定理不仅适用于二维平面，也适用于三维空间中的直角坐标系中的点到原点的距离关系。它是许多几何和代数问题的基础，同时也是现代数学教育中的一个重要概念。', response_metadata={'model_name': 'qwen-turbo', 'finish_reason': 'stop', 'request_id': '09f736aa-84b9-9a31-b249-1cf1a779b034', 'token_usage': {'input_tokens': 22, 'output_tokens': 142, 'total_tokens': 164}}, id='run-0978a38d-3611-4c89-97e2-05280f578591-0')"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_core.messages import HumanMessage, SystemMessage\n",
    "\n",
    "messages = [\n",
    "    SystemMessage(content=\"你是一个数学专家\"),\n",
    "    HumanMessage(content=\"什么是勾股定理\"),\n",
    "]\n",
    "\n",
    "chat.invoke(messages)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 流式模式"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "勾|股|定|理是古希腊数学|家毕达哥拉斯发现的一个几何|学基本定理，也被称为毕|达哥拉斯定理。它描述|了直角三角形中各边|之间的关系：在一个直角三角形|中，直角边（即与|直角相邻的两边）的平方|和等于斜边（即三角形|最长边，对直角）的|平方。用数学公式表示就是：\n",
      "\n",
      "|如果直角三角形的两条直|角边长分别为a和b，|斜边长为c，那么有|：\n",
      "a² + b² = c|²\n",
      "\n",
      "这个定理不仅适用于二维|平面，其抽象的概念在三维空间|和其他高维空间中也有类似的表述|。它是解决涉及直角三角形|问题的关键工具，在数学、物理学、|工程学等多个领域都有广泛的应用。||"
     ]
    }
   ],
   "source": [
    "for chunk in chat.stream(messages):\n",
    "    print(chunk.content, end=\"|\", flush=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 缓存"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.globals import set_llm_cache\n",
    "from langchain.cache import InMemoryCache"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 61.1 ms, sys: 15.2 ms, total: 76.3 ms\n",
      "Wall time: 818 ms\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "AIMessage(content='小王剪了一个中分，然后他就变成了小全。', response_metadata={'model_name': 'qwen-turbo', 'finish_reason': 'stop', 'request_id': '11862524-8349-94b5-9a9a-456be26f7336', 'token_usage': {'input_tokens': 11, 'output_tokens': 13, 'total_tokens': 24}}, id='run-764a7128-884a-45c8-8316-80f3d312a919-0')"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "set_llm_cache(InMemoryCache())\n",
    "\n",
    "# The first time, it is not yet in cache, so it should take longer\n",
    "chat.invoke(\"说一个笑话\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 827 µs, sys: 106 µs, total: 933 µs\n",
      "Wall time: 982 µs\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "AIMessage(content='小王剪了一个中分，然后他就变成了小全。', response_metadata={'model_name': 'qwen-turbo', 'finish_reason': 'stop', 'request_id': '11862524-8349-94b5-9a9a-456be26f7336', 'token_usage': {'input_tokens': 11, 'output_tokens': 13, 'total_tokens': 24}}, id='run-764a7128-884a-45c8-8316-80f3d312a919-0')"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "# The second time it is, so it goes faster\n",
    "chat.invoke(\"说一个笑话\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 自定义Chat model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Any, AsyncIterator, Dict, Iterator, List, Optional\n",
    "\n",
    "from langchain_core.callbacks import (\n",
    "    AsyncCallbackManagerForLLMRun,\n",
    "    CallbackManagerForLLMRun,\n",
    ")\n",
    "from langchain_core.language_models import BaseChatModel, SimpleChatModel\n",
    "from langchain_core.messages import AIMessageChunk, BaseMessage, HumanMessage\n",
    "from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult\n",
    "from langchain_core.runnables import run_in_executor\n",
    "\n",
    "\n",
    "class CustomChatModelAdvanced(BaseChatModel):\n",
    "    \"\"\"A custom chat model that echoes the first `n` characters of the input.\n",
    "\n",
    "    When contributing an implementation to LangChain, carefully document\n",
    "    the model including the initialization parameters, include\n",
    "    an example of how to initialize the model and include any relevant\n",
    "    links to the underlying models documentation or API.\n",
    "\n",
    "    Example:\n",
    "\n",
    "        .. code-block:: python\n",
    "\n",
    "            model = CustomChatModel(n=2)\n",
    "            result = model.invoke([HumanMessage(content=\"hello\")])\n",
    "            result = model.batch([[HumanMessage(content=\"hello\")],\n",
    "                                 [HumanMessage(content=\"world\")]])\n",
    "    \"\"\"\n",
    "\n",
    "    model_name: str\n",
    "    \"\"\"The name of the model\"\"\"\n",
    "    n: int\n",
    "    \"\"\"The number of characters from the last message of the prompt to be echoed.\"\"\"\n",
    "\n",
    "    def _generate(\n",
    "        self,\n",
    "        messages: List[BaseMessage],\n",
    "        stop: Optional[List[str]] = None,\n",
    "        run_manager: Optional[CallbackManagerForLLMRun] = None,\n",
    "        **kwargs: Any,\n",
    "    ) -> ChatResult:\n",
    "        \"\"\"Override the _generate method to implement the chat model logic.\n",
    "\n",
    "        This can be a call to an API, a call to a local model, or any other\n",
    "        implementation that generates a response to the input prompt.\n",
    "\n",
    "        Args:\n",
    "            messages: the prompt composed of a list of messages.\n",
    "            stop: a list of strings on which the model should stop generating.\n",
    "                  If generation stops due to a stop token, the stop token itself\n",
    "                  SHOULD BE INCLUDED as part of the output. This is not enforced\n",
    "                  across models right now, but it's a good practice to follow since\n",
    "                  it makes it much easier to parse the output of the model\n",
    "                  downstream and understand why generation stopped.\n",
    "            run_manager: A run manager with callbacks for the LLM.\n",
    "        \"\"\"\n",
    "        # Replace this with actual logic to generate a response from a list\n",
    "        # of messages.\n",
    "        last_message = messages[-1]\n",
    "        tokens = last_message.content[: self.n]\n",
    "        message = AIMessage(\n",
    "            content=tokens,\n",
    "            additional_kwargs={},  # Used to add additional payload (e.g., function calling request)\n",
    "            response_metadata={  # Use for response metadata\n",
    "                \"time_in_seconds\": 3,\n",
    "            },\n",
    "        )\n",
    "        ##\n",
    "\n",
    "        generation = ChatGeneration(message=message)\n",
    "        return ChatResult(generations=[generation])\n",
    "\n",
    "    def _stream(\n",
    "        self,\n",
    "        messages: List[BaseMessage],\n",
    "        stop: Optional[List[str]] = None,\n",
    "        run_manager: Optional[CallbackManagerForLLMRun] = None,\n",
    "        **kwargs: Any,\n",
    "    ) -> Iterator[ChatGenerationChunk]:\n",
    "        \"\"\"Stream the output of the model.\n",
    "\n",
    "        This method should be implemented if the model can generate output\n",
    "        in a streaming fashion. If the model does not support streaming,\n",
    "        do not implement it. In that case streaming requests will be automatically\n",
    "        handled by the _generate method.\n",
    "\n",
    "        Args:\n",
    "            messages: the prompt composed of a list of messages.\n",
    "            stop: a list of strings on which the model should stop generating.\n",
    "                  If generation stops due to a stop token, the stop token itself\n",
    "                  SHOULD BE INCLUDED as part of the output. This is not enforced\n",
    "                  across models right now, but it's a good practice to follow since\n",
    "                  it makes it much easier to parse the output of the model\n",
    "                  downstream and understand why generation stopped.\n",
    "            run_manager: A run manager with callbacks for the LLM.\n",
    "        \"\"\"\n",
    "        last_message = messages[-1]\n",
    "        tokens = last_message.content[: self.n]\n",
    "\n",
    "        for token in tokens:\n",
    "            chunk = ChatGenerationChunk(message=AIMessageChunk(content=token))\n",
    "\n",
    "            if run_manager:\n",
    "                # This is optional in newer versions of LangChain\n",
    "                # The on_llm_new_token will be called automatically\n",
    "                run_manager.on_llm_new_token(token, chunk=chunk)\n",
    "\n",
    "            yield chunk\n",
    "\n",
    "        # Let's add some other information (e.g., response metadata)\n",
    "        chunk = ChatGenerationChunk(\n",
    "            message=AIMessageChunk(content=\"\", response_metadata={\"time_in_sec\": 3})\n",
    "        )\n",
    "        if run_manager:\n",
    "            # This is optional in newer versions of LangChain\n",
    "            # The on_llm_new_token will be called automatically\n",
    "            run_manager.on_llm_new_token(token, chunk=chunk)\n",
    "        yield chunk\n",
    "\n",
    "    @property\n",
    "    def _llm_type(self) -> str:\n",
    "        \"\"\"Get the type of language model used by this chat model.\"\"\"\n",
    "        return \"echoing-chat-model-advanced\"\n",
    "\n",
    "    @property\n",
    "    def _identifying_params(self) -> Dict[str, Any]:\n",
    "        \"\"\"Return a dictionary of identifying parameters.\n",
    "\n",
    "        This information is used by the LangChain callback system, which\n",
    "        is used for tracing purposes make it possible to monitor LLMs.\n",
    "        \"\"\"\n",
    "        return {\n",
    "            # The model name allows users to specify custom token counting\n",
    "            # rules in LLM monitoring applications (e.g., in LangSmith users\n",
    "            # can provide per token pricing for their model and monitor\n",
    "            # costs for the given LLM.)\n",
    "            \"model_name\": self.model_name,\n",
    "        }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content='Meo', response_metadata={'time_in_seconds': 3}, id='run-38fb6c01-bc91-4b3a-84b5-989a52a0a245-0')"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_core.messages import (\n",
    "    AIMessage,\n",
    "    BaseMessage,\n",
    "    FunctionMessage,\n",
    "    HumanMessage,\n",
    "    SystemMessage,\n",
    "    ToolMessage,\n",
    ")\n",
    "\n",
    "model = CustomChatModelAdvanced(n=3, model_name=\"my_custom_model\")\n",
    "\n",
    "model.invoke(\n",
    "    [\n",
    "        HumanMessage(content=\"hello!\"),\n",
    "        AIMessage(content=\"Hi there human!\"),\n",
    "        HumanMessage(content=\"Meow!\"),\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "c|a|t||"
     ]
    }
   ],
   "source": [
    "# 输入也支持字符串，可以等同于`[HumanMessage(content=\"cat vs dog\")]`\n",
    "for chunk in model.stream(\"cat vs dog\"):\n",
    "    print(chunk.content, end=\"|\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 记录消耗tokens"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tokens Used: 153\n",
      "\tPrompt Tokens: 22\n",
      "\tCompletion Tokens: 131\n",
      "Successful Requests: 1\n",
      "Total Cost (CYN): ¥0.0027960000000000003\n"
     ]
    }
   ],
   "source": [
    "from langchain_core.messages import HumanMessage, SystemMessage\n",
    "from callbacks.manager import get_generic_llms_callback\n",
    "\n",
    "messages = [\n",
    "        SystemMessage(content=\"你是一个数学专家\"),\n",
    "        HumanMessage(content=\"什么是勾股定理\"),\n",
    "    ]\n",
    "\n",
    "with get_generic_llms_callback() as cb:\n",
    "    chat.invoke(messages)\n",
    "    print(cb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "勾|股|定|理是古希腊数学|家毕达哥拉斯发现的一个几何|学基本定理，也被称为毕|达哥拉斯定理。它描述|了直角三角形中各边|之间的关系：在一个直角三角形|中，直角边（即与|直角相邻的两边）的平方|和等于斜边（即三角形|最长边，对直角）的|平方。用数学公式表示就是：\n",
      "\n",
      "|如果直角三角形的两条直|角边长分别为a和b，|斜边长为c，那么有|：\n",
      "a² + b² = c|²\n",
      "\n",
      "这个定理不仅适用于二维|平面，其抽象的概念在三维空间|和其他高维空间中也有类似的表述|。它是解决涉及直角三角形|问题的关键工具，在数学、物理学、|工程学等多个领域都有广泛的应用。|||\n",
      "Tokens Used: 190\n",
      "\tPrompt Tokens: 22\n",
      "\tCompletion Tokens: 168\n",
      "Successful Requests: 1\n",
      "Total Cost (CYN): ¥0.003536\n"
     ]
    }
   ],
   "source": [
    "from tongyi.chat_model import CustomChatTongyi\n",
    "\n",
    "# dashscope_api_key作为参数传入\n",
    "# 或者配置环境变量`DASHSCOPE_API_KEY`\n",
    "chat = CustomChatTongyi()\n",
    "\n",
    "with get_generic_llms_callback() as cb:\n",
    "    for chunk in chat.stream(messages):\n",
    "        print(chunk.content, end=\"|\", flush=True)\n",
    "    \n",
    "    print()\n",
    "    print(cb)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "langchain",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
